import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from mpl_toolkits.mplot3d import Axes3D
import random


def is_convergent(old_parameters, new_parameters, w_threshold=0.1, mean_threshold=1, cov_threshold=1):
    r_w1 = abs(old_parameters['w1'] - new_parameters['w1'])
    r_w2 = abs(old_parameters['w2'] - new_parameters['w2'])
    r_w3 = abs(old_parameters['w3'] - new_parameters['w3'])
    r_mean1 = np.linalg.norm(old_parameters['mean1'] - new_parameters['mean1'], ord=2)
    r_mean2 = np.linalg.norm(old_parameters['mean2'] - new_parameters['mean2'], ord=2)
    r_mean3 = np.linalg.norm(old_parameters['mean3'] - new_parameters['mean3'], ord=2)
    r_cov1 = np.linalg.norm(old_parameters['cov1'] - new_parameters['cov1'])
    r_cov2 = np.linalg.norm(old_parameters['cov2'] - new_parameters['cov2'])
    r_cov3 = np.linalg.norm(old_parameters['cov3'] - new_parameters['cov3'])
    print(r_w1, r_w2, r_w3, r_mean1, r_mean2, r_mean3, r_cov1, r_cov2, r_cov3)
    if r_w1 < w_threshold and r_w2 < w_threshold and r_w3 < w_threshold and \
            r_mean1 < mean_threshold and r_mean2 < mean_threshold and r_mean3 < mean_threshold and \
            r_cov1 < cov_threshold and r_cov2 < cov_threshold and r_cov3 < cov_threshold:
        return True
    else:
        return False


def update_cov(p, mean):
    s = np.zeros((2, 2))
    for i in range(n):
        temp = (data[i] - mean).reshape(2, 1)
        s += np.matmul(temp, temp.T) * p[i]
    cov = s / np.sum(p)
    return cov


def update_mean(p):
    s = np.zeros(2)
    for i in range(n):
        s += data[i] * p[i]
    mean = s / np.sum(p)
    return mean


def draw_points(index):
    points1 = data[index == 0, :]
    points2 = data[index == 1, :]
    points3 = data[index == 2, :]
    x, y = points1.T
    plt.scatter(x, y)
    x, y = points2.T
    plt.scatter(x, y)
    x, y = points3.T
    plt.scatter(x, y)
    plt.show()


def draw_Gaussian(Gaussian1, Gaussian2, Gaussian3, w1, w2, w3):
    # 生成二维网格平面
    X, Y = np.meshgrid(np.linspace(-100, 100, 1000), np.linspace(-100, 100, 1000))
    # 二维坐标数据
    d = np.dstack([X, Y])
    Z1 = Gaussian1.pdf(d).reshape(1000, 1000) * w1
    Z2 = Gaussian2.pdf(d).reshape(1000, 1000) * w2
    Z3 = Gaussian3.pdf(d).reshape(1000, 1000) * w3
    Z = Z1 + Z2 + Z3
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='seismic', alpha=0.8)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.show()


def iris():
    iris = []
    file = open('iris1.txt')
    for line in file:
        temp = []
        data = line.split('\t', -1)
        temp.append(data[2])
        temp.append(data[3])
        iris.append(temp)
    iris = np.array(iris)
    return iris

data=iris()

# mean_1, cov_1 = np.array([20, 30]), np.array([[80, 10], [10, 85]])
# mean_2, cov_2 = np.array([-40, 7]), np.array([[77, 12], [12, 74]])
# mean_3, cov_3 = np.array([-2, -55]), np.array([[86, 11], [11, 81]])
# data1 = np.random.multivariate_normal(mean_1, cov_1, 500)
# data2 = np.random.multivariate_normal(mean_2, cov_2, 800)
# data3 = np.random.multivariate_normal(mean_3, cov_3, 300)
# data = np.vstack((data1, data2, data3))
n = len(data)

w1, mean1, cov1 = 0.2, np.random.random(2) * 200 - 100, np.eye(2) * random.randrange(1, 100)
w2, mean2, cov2 = 0.3, np.random.random(2) * 200 - 100, np.eye(2) * random.randrange(1, 100)
w3, mean3, cov3 = 0.5, np.random.random(2) * 200 - 100, np.eye(2) * random.randrange(1, 100)

Gaussian1 = multivariate_normal(mean=mean1, cov=cov1)
Gaussian2 = multivariate_normal(mean=mean2, cov=cov2)
Gaussian3 = multivariate_normal(mean=mean3, cov=cov3)

old_parameters = {
    'w1': w1,
    'w2': w2,
    'w3': w3,
    'mean1': mean1,
    'mean2': mean2,
    'mean3': mean3,
    'cov1': cov1,
    'cov2': cov2,
    'cov3': cov3,
}
new_parameters = {
    'w1': w1,
    'w2': w2,
    'w3': w3,
    'mean1': mean1,
    'mean2': mean2,
    'mean3': mean3,
    'cov1': cov1,
    'cov2': cov2,
    'cov3': cov3,
}



first = True
while (is_convergent(old_parameters, new_parameters) == False or first):
    first = False
    out1 = Gaussian1.pdf(data)
    out2 = Gaussian2.pdf(data)
    out3 = Gaussian3.pdf(data)
    total = w1 * out1 + w2 * out2 + w3 * out3
    p1 = (w1 * out1) / total
    p2 = (w2 * out2) / total
    p3 = (w3 * out3) / total
    index = np.vstack((p1, p2, p3)).argmax(axis=0)
    draw_points(index)

    old_parameters['w1'], old_parameters['w2'], old_parameters['w3'] = w1, w2, w3
    old_parameters['mean1'], old_parameters['mean2'], old_parameters['mean3'] = mean1, mean2, mean3
    old_parameters['cov1'], old_parameters['cov2'], old_parameters['cov3'] = cov1, cov2, cov3
    w1 = np.sum(p1) / n
    w2 = np.sum(p2) / n
    w3 = np.sum(p3) / n
    mean1 = update_mean(p1)
    mean2 = update_mean(p2)
    mean3 = update_mean(p3)
    cov1 = update_cov(p1, mean1)
    cov2 = update_cov(p2, mean2)
    cov3 = update_cov(p3, mean3)
    Gaussian1 = multivariate_normal(mean=mean1, cov=cov1)
    Gaussian2 = multivariate_normal(mean=mean2, cov=cov2)
    Gaussian3 = multivariate_normal(mean=mean3, cov=cov3)
    draw_Gaussian(Gaussian1, Gaussian2, Gaussian3, w1, w2, w3)

    new_parameters['w1'], new_parameters['w2'], new_parameters['w3'] = w1, w2, w3
    new_parameters['mean1'], new_parameters['mean2'], new_parameters['mean3'] = mean1, mean2, mean3
    new_parameters['cov1'], new_parameters['cov2'], new_parameters['cov3'] = cov1, cov2, cov3
